import asyncio
import json
import time

import ray
from datasets import load_dataset

from lagent.distributed.ray_serve import AsyncAgentRayActor
from lagent.llms import INTERNLM2_META
from lagent.llms.lmdeploy_wrapper import AsyncLMDeployPipeline

# 初始化 Ray 分布式框架
ray.init()

# 设置异步事件循环
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)

# 配置 LMDeploy 模型参数
model = dict(
    type=AsyncLMDeployPipeline,
    path='internlm/internlm2_5-7b-chat',  # 模型路径
    meta_template=INTERNLM2_META,  # 使用 InternLM2 的元模板
    tp=1,  # 张量并行度
    top_k=1,  # 只保留概率最高的候选词
    temperature=1.0,  # 采样温度
    stop_words=['<|im_end|>', '<|action_end|>'],  # 停止词
    max_new_tokens=1024,  # 最大生成长度
)

# ----------------------- 测试数学解题功能 -----------------------
print('-' * 80, 'interpreter', '-' * 80)

# 加载数学问题数据集
ds = load_dataset('lighteval/MATH', split='test')
problems = [item['problem'] for item in ds.select(range(5000))]  # 选取5000个问题

# 配置数学解题代理
coder = dict(
    type='lagent.agents.stream.AsyncMathCoder',  # 使用异步数学解题代理
    llm=model,  # 使用上面配置的模型
    interpreter=dict(
        type='AsyncIPythonInterpreter',  # 使用异步 Python 解释器
        max_kernels=300  # 最大并发内核数
    ),
)

# 开始批量解题
tic = time.time()  # 记录开始时间

# 创建两个 Ray Actor，每个使用一个 GPU
actor1 = AsyncAgentRayActor(coder.copy(), num_gpus=1)
actor2 = AsyncAgentRayActor(coder.copy(), num_gpus=1)

# 将问题平均分配给两个 Actor
corots = [
    actor1(query, session_id=i)
    for i, query in enumerate(problems[:len(problems) // 2])  # 前半部分问题
]
corots += [
    actor2(query, session_id=i)
    for i, query in enumerate(problems[len(problems) // 2:])  # 后半部分问题
]

# 并发执行所有任务
results = loop.run_until_complete(asyncio.gather(*corots))

print('-' * 120)
print(f'time elapsed: {time.time() - tic}')  # 打印总耗时

# 收集两个 Actor 的解题步骤
all_step = ray.get([
    actor1.agent_actor.get_steps.remote(i) for i in range(len(problems) // 2)
])
all_step += ray.get([
    actor2.agent_actor.get_steps.remote(i)
    for i in range(len(problems[len(problems) // 2:]))
])

# 保存所有解题步骤到文件
with open('./tmp_1.json', 'w') as f:
    json.dump(all_step, f, ensure_ascii=False, indent=4)
